import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import matplotlib.pyplot as plt, seaborn as sns
from sklearn.preprocessing import RobustScaler
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, recall_score, f1_score
from sklearn.metrics import confusion_matrix, classification_report
from imblearn.over_sampling import SMOTE
import time

# 记录开始时间
start = time.time()

# 读取数据集
ds = pd.read_csv('./dataset/creditcard.csv')

# 缩放时间和金额
rob_scaler = RobustScaler()
ds['Amount'] = rob_scaler.fit_transform(ds['Amount'].values.reshape(-1, 1))
ds['Time'] = rob_scaler.fit_transform(ds['Time'].values.reshape(-1, 1))

# 准备特征和标签
X = ds.iloc[:, :-1]
y = ds.iloc[:, -1]

# 数据预处理和拆分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=3, train_size=.9, stratify=y)

# 使用SMOTE处理数据不平衡
smote = SMOTE(random_state=3)
X_train_resampled, y_train_resampled = smote.fit_resample(X_train, y_train)
print(X_train_resampled.shape)

# 转换为PyTorch张量
X_train_tensor = torch.tensor(X_train_resampled.values, dtype=torch.float32)
y_train_tensor = torch.tensor(y_train_resampled.values, dtype=torch.float32)
X_test_tensor = torch.tensor(X_test.values, dtype=torch.float32)
y_test_tensor = torch.tensor(y_test.values, dtype=torch.float32)

# model training and prediction
clf = RandomForestClassifier(n_estimators=100, criterion='gini')
clf.fit(X_train_tensor, y_train_tensor)
y_pred = clf.predict(X_test_tensor)
end_time = time.time() - start

# acc data
accuracy = accuracy_score(y_test_tensor, y_pred)
print(f"RandomForest 准确率:{accuracy * 100:.3f}%")
recall = recall_score(y_test_tensor, y_pred, average='macro')
print(f"RandomForest 召回率:{recall * 100:.3f}%")
f1 = f1_score(y_test, y_pred, average='macro')
print(f"RandomForest F1:{f1 * 100:.3f}%")
print(f"used time:{end_time}")

report = classification_report(y_test_tensor, y_pred)
cm = confusion_matrix(y_test_tensor, y_pred)
print("RandomForest Classification Report:\n")
print(report)
print("RandomForest Confusion Matrix:\n")
print(cm)

# plot heatmap for res visualization
cm = pd.crosstab(y_test_tensor, y_pred, rownames=['Actual'], colnames=['Predicted'])
fig, (ax1) = plt.subplots(ncols=1, figsize=(5,5))
sns.heatmap(cm, 
            xticklabels=['Not Fraud', 'Fraud'],
            yticklabels=['Not Fraud', 'Fraud'],
            annot=True,ax=ax1,
            linewidths=.2,linecolor="Darkblue", cmap="Blues")
plt.title('Confusion Matrix', fontsize=14)
plt.show()